[TRTLLM-15177][chore] Kimi K3 post-merge cleanup: config/import/test hygiene + L0 wiring - #17413

Summary

Post-merge cleanup follow-ups deferred from #17269 (Kimi K3 / KimiLinear model), tracked in TRTLLM-15177. Low-risk, K3-scoped hygiene.

Included

• config: KimiLinearConfig.keys_to_ignore_at_inference list to tuple (immutable class-level default).
• modeling_kimi_linear.py: hoist gc/json/ExitStack/safetensors.safe_open to module level (removed 3 redundant local import gc), drop the json->_json alias, move the latent_moe_use_norm assert to the top of KimiK3MoERuntime.__init__.
• modeling_kimi_linear.py refactors (behavior-neutral, code moved verbatim):
  • split the ~450-line KimiLinearForCausalLM.load_weights into focused methods (_validate_checkpoint_keys, _load_trunk_params, _load_expert_slices, _finalize_weight_load) behind a short orchestrator;
  • extract the duplicated FP8 weight-read module swap (nested _swap closures in the MoE-MLP and MLA converters plus the inline KDA o_proj conversion) into a module-level _swap_linear_to_fp8_weight_read helper.
• Move the test-only reference MoE block (KimiK3SparseMoeBlock + _moe_kernels/_mxfp4 helpers, used only by test_kimi_k3_situ_moe.py) out of the runtime package to tests/unittest/_torch/modules/moe/kimi_k3_ref_moe/, rewriting relative imports to absolute. The runtime package keeps the gate and shared MLP/RMSNorm pieces modeling_kimi_linear.py uses.
• KDA fused-verify parity test: real KimiLinearConfig instead of SimpleNamespace, @torch.no_grad() decorators, module-level imports, drop the __main__ block, ruff-format.
• L0: wire 6 K3 unit suites into l0_b200.yml; the 7th (test_kda_prefill_op.py) is deliberately excluded pending the B200 prefill-op NaN (TRTLLM-15204).

Deferred (still tracked in TRTLLM-15177)

• Folding kimi_k3_mla into modeling_kimi_linear.py: small on this base (one ~260-line module with a single importer, already built on the general MLA class), but kimi_k3_mla_attention.py has diverged on the K3 feature branch (decode-backend selection changes not yet in main), so deleting the module here would turn the next branch sync into modify/delete conflicts. Deferred until the branches converge.
• Migrating kimi_k3_moe into modeling_kimi_linear.py: [None][perf] Shard Kimi K3 MLPs with tensor parallelism #17312 simplifies the MoE package (standard GatedMLP for the shared experts and first dense layer) and is the right base for this migration; it targets the K3 feature branch and is not in main yet. Deferred until [None][perf] Shard Kimi K3 MLPs with tensor parallelism #17312 reaches main.
• Overlap note: the reference-MoE relocation in this PR and [None][perf] Shard Kimi K3 MLPs with tensor parallelism #17312 move the same test-only code into the unit-test tree under different names (kimi_k3_ref_moe/ package here vs flat *_test_utils.py modules there). The PRs target different branches (main vs feat/kimi_k3), so neither blocks the other; the branch-sync merge after both land should keep [None][perf] Shard Kimi K3 MLPs with tensor parallelism #17312's placement and drop the duplicate.
• Removing the communication_method plumbing was intentionally dropped. It is not dead: test_kimi_k3_situ_moe.py passes communication_method="ALLGATHER", and communication_factory.py uses it as the default when TRTLLM_FORCE_COMM_METHOD is unset.

Validation

Validated on a Blackwell node (GB300, SM103 — representative of B200; not run on B200 hardware itself):

• All 6 K3 unit suites wired into l0_b200.yml pass, including the changed test_kimi_kda_fused_verify_parity.py. The excluded test_kda_prefill_op.py fails deterministically on B200 (NaN in every compute-running case, TRTLLM-15204); the full set of 7 passed on GB300 out-of-CI.
• Test collection (pytest --collect-only) of all 7 suites (6 wired + 1 excluded) is clean, confirming the module-level modeling_kimi_linear import is safe at collection time; imports of the relocated reference MoE module verified.
• Weight-load refactor checked against a real checkpoint: a layer-truncated Kimi K3 TP4 load + generation was run on the pre-refactor and post-refactor trees; generated token ids match exactly and per-token logprobs are bit-identical across the two runs.
• pre-commit clean on all changed files.

Related: TRTLLM-15177

Dev Engineer Review

• Changed KimiLinearConfig.keys_to_ignore_at_inference to an immutable tuple.
• Refactored Kimi K3 checkpoint loading and FP8 conversion.
• Added validation, typed helpers, streamed expert loading, page eviction, and expert-slot completeness checks.
• Changed latent-MoE validation from assert to ValueError.
• Moved the HF-parity reference MoE implementation to the test package.
• Removed test-only MoE symbols from runtime exports.
• Deferred kimi_k3_mla and kimi_k3_moe consolidation.
• Kept communication_method plumbing unchanged.
• TP4 checkpoint generation produced identical token IDs and per-token logprobs before and after the refactor.

QA Engineer Review

• Modified tests/integration/test_lists/test-db/l0_b200.yml.
• Added these L0 entries:
  • test_kimi_kda_fused_verify_parity.py
  • test_kimi_kda_verify_parity.py
  • test_kda_cache_soundness.py
  • test_kda_prefill_state_parity.py
  • test_attn_res_op.py
  • test_kimi_k3_situ_moe.py
• Excluded test_kda_prefill_op.py because its compute cases produce NaNs on B200. The exclusion references TRTLLM-15204.
• Updated test_fused_vs_sequential_two_rounds.
• Updated Kimi K3 SiTU MoE tests, including fused/reference parity, mutation, error handling, mapping, split selection, and TP-shard loader tests.
• The modified test files are covered by test-db/l0_b200.yml.
• Test collection and pre-commit checks passed.
• Two CI runs failed. A later fail-fast-disabled run passed after automatic retries.
• Validation was performed on GB300.
• Verdict: sufficient.